import os
import uuid
import json
import shutil
from flask import Flask, render_template, request, jsonify, send_from_directory, send_file, redirect, url_for
from werkzeug.utils import secure_filename
from modules.llm_processor import LLMProcessor
from modules.image_generator import ImageGenerator
from modules.style_enhancer import StyleEnhancer
from modules.poster_creator import PosterCreator
from modules.clip_filter import CLIPFilter
import config

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# Initialize components
llm_processor = LLMProcessor()
image_generator = ImageGenerator()
style_enhancer = StyleEnhancer()
clip_filter = CLIPFilter()
poster_creator = PosterCreator()


@app.route('/')
def index():
    return render_template('index.html', styles=config.STYLE_TEMPLATES.keys())


@app.route('/generate', methods=['POST'])
def generate_poster():
    # Get paper information
    paper_title = request.form.get('title')
    paper_abstract = request.form.get('abstract')
    paper_authors = request.form.get('authors')
    selected_style = request.form.get('style')
    custom_style = request.form.get('custom_style')

    # Handle paper PDF if uploaded
    paper_file = None
    if 'paper_file' in request.files:
        file = request.files['paper_file']
        if file.filename != '':
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(file_path)
            paper_file = file_path

    # Generate unique ID for this request
    request_id = str(uuid.uuid4())

    # Process using our pipeline
    try:
        # 1. Use LLM to process paper content and create structured poster content
        poster_content = llm_processor.process_paper(
            title=paper_title,
            abstract=paper_abstract,
            authors=paper_authors,
            paper_file=paper_file
        )

        # 2. Generate images based on content
        # 直接传递用户自定义风格到图像生成器
        style_for_image = custom_style if selected_style == 'custom' else None
        image_paths = image_generator.generate(poster_content, style_for_image)

        # 3. Apply style enhancement
        style_description = custom_style if selected_style == 'custom' else selected_style
        enhanced_content = style_enhancer.enhance(poster_content, style_description)

        # 4. Filter images if multiple were generated
        if len(image_paths) > 1:
            image_paths = clip_filter.filter_images(image_paths, enhanced_content, style_description)

        # 5. Create the final poster
        poster_path = poster_creator.create_poster(enhanced_content, image_paths, style_description, request_id)

        # 获取生成的图像路径（用于直接下载）
        background_image_path = None
        if image_paths and len(image_paths) > 0:
            background_image_path = image_paths[0]

        # 确保输出目录存在
        os.makedirs(config.OUTPUT_DIR, exist_ok=True)

        # 6. Return results
        preview_filename = f"{request_id}_preview.png"
        poster_filename = f"{request_id}.pdf"
        bg_filename = f"background_{request_id}.jpg"
        
        preview_url = f"/output/{preview_filename}"
        poster_url = f"/output/{poster_filename}"
        
        # 复制背景图像到输出目录并设置正确的URL
        background_url = None
        if background_image_path and os.path.exists(background_image_path):
            background_output_path = os.path.join(config.OUTPUT_DIR, bg_filename)
            try:
                shutil.copy2(background_image_path, background_output_path)
                background_url = f"/output/{bg_filename}"
                print(f"背景图片已复制到: {background_output_path}")
            except Exception as e:
                print(f"复制背景图片时出错: {str(e)}")
        
        # 检查生成的文件是否存在
        preview_path = os.path.join(config.OUTPUT_DIR, preview_filename)
        poster_path = os.path.join(config.OUTPUT_DIR, poster_filename)
        
        if not os.path.exists(preview_path):
            print(f"警告: 预览图片不存在: {preview_path}")
        
        if not os.path.exists(poster_path):
            print(f"警告: PDF海报不存在: {poster_path}")
        
        # 始终返回JSON响应
        response_data = {
            'success': True,
            'preview_url': preview_url,
            'poster_url': poster_url,
            'background_url': background_url,
            'title': paper_title
        }
        
        print(f"返回响应: {json.dumps(response_data)}")
        return jsonify(response_data)

    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"生成海报时出错: {str(e)}\n{error_trace}")
        return jsonify({'success': False, 'error': str(e)})


# 删除这个路由，使用统一的 /output/<filename> 路由
# @app.route('/posters/<filename>')
# def get_poster(filename):
#     return send_file(os.path.join(config.OUTPUT_DIR, filename))

@app.route('/output/<filename>')
def get_output(filename):
    """查看输出文件的路由"""
    file_path = os.path.join(config.OUTPUT_DIR, filename)
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        return "文件不存在", 404
    
    # 如果是图片文件，重定向到查看图片页面
    if filename.endswith(('.jpg', '.jpeg', '.png')):
        return redirect(url_for('view_image', filename=filename))
    
    try:
        # 根据文件类型设置正确的 MIME 类型
        mimetype = None
        if filename.endswith('.pdf'):
            mimetype = 'application/pdf'
        
        return send_from_directory(
            config.OUTPUT_DIR, 
            filename, 
            mimetype=mimetype,
            as_attachment=False
        )
    except Exception as e:
        print(f"发送文件时出错: {str(e)}")
        return f"发送文件时出错: {str(e)}", 500

@app.route('/result/<request_id>')
def show_result(request_id):
    # 构建URL
    preview_url = f"/output/{request_id}_preview.png"
    poster_url = f"/output/{request_id}.pdf"
    bg_filename = f"background_{request_id}.jpg"
    background_url = f"/output/{bg_filename}"
    
    # 获取标题（可选，如果需要的话）
    title = "生成的海报"  # 默认标题
    
    return render_template('result.html', 
                          title=title,
                          preview_url=preview_url,
                          poster_url=poster_url,
                          background_url=background_url)


@app.route('/download/<filename>')
def download_file(filename):
    """专用于下载文件的路由，强制浏览器下载而不是打开"""
    file_path = os.path.join(config.OUTPUT_DIR, filename)
    if not os.path.exists(file_path):
        print(f"下载文件不存在: {file_path}")
        return "文件不存在", 404
    
    try:
        # 设置正确的 MIME 类型
        mimetype = None
        if filename.endswith('.pdf'):
            mimetype = 'application/pdf'
        elif filename.endswith('.jpg') or filename.endswith('.jpeg'):
            mimetype = 'image/jpeg'
        elif filename.endswith('.png'):
            mimetype = 'image/png'
        
        return send_from_directory(
            config.OUTPUT_DIR, 
            filename, 
            mimetype=mimetype,
            as_attachment=True  # 强制下载
        )
    except Exception as e:
        print(f"下载文件时出错: {str(e)}")
        return f"下载文件时出错: {str(e)}", 500


@app.route('/view/<filename>')
def view_image(filename):
    """查看图片的路由"""
    file_path = os.path.join(config.OUTPUT_DIR, filename)
    if not os.path.exists(file_path):
        return "文件不存在", 404
    
    return render_template('image_view.html', filename=filename)


if __name__ == '__main__':
    app.run(debug=True)